{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1: Getting started"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.utils import to_categorical\n",
    "from sklearn.datasets import fetch_openml\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder, StandardScaler\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "seed = 0\n",
    "np.random.seed(seed)\n",
    "import tensorflow as tf\n",
    "tf.random.set_seed(seed)\n",
    "import os\n",
    "os.environ['PATH'] = '/opt/Xilinx/Vivado/2019.2/bin:' + os.environ['PATH']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fetch the jet tagging dataset from Open ML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = fetch_openml('hls4ml_lhc_jets_hlf')\n",
    "X, y = data['data'], data['target']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's print some information about the dataset\n",
    "Print the feature names and the dataset shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(data['feature_names'])\n",
    "print(X.shape, y.shape)\n",
    "print(X[:5])\n",
    "print(y[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you saw above, the `y` target is an array of strings, e.g. \\['g', 'w',...\\] etc.\n",
    "We need to make this a \"One Hot\" encoding for the training.\n",
    "Then, split the dataset into training and validation sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "le = LabelEncoder()\n",
    "y = le.fit_transform(y)\n",
    "y = to_categorical(y, 5)\n",
    "X_train_val, X_test, y_train_val, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n",
    "print(y[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler = StandardScaler()\n",
    "X_train_val = scaler.fit_transform(X_train_val)\n",
    "X_test = scaler.transform(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('X_train_val.npy', X_train_val)\n",
    "np.save('X_test.npy', X_test)\n",
    "np.save('y_train_val.npy', y_train_val)\n",
    "np.save('y_test.npy', y_test)\n",
    "np.save('classes.npy', le.classes_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now construct a model\n",
    "We'll use 3 hidden layers with 64, then 32, then 32 neurons. Each layer will use `relu` activation.\n",
    "Add an output layer with 5 neurons (one for each class), then finish with Softmax activation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense, Activation, BatchNormalization\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "from tensorflow.keras.regularizers import l1\n",
    "from callbacks import all_callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Sequential()\n",
    "model.add(Dense(64, input_shape=(16,), name='fc1', kernel_initializer='lecun_uniform', kernel_regularizer=l1(0.0001)))\n",
    "model.add(Activation(activation='relu', name='relu1'))\n",
    "model.add(Dense(32, name='fc2', kernel_initializer='lecun_uniform', kernel_regularizer=l1(0.0001)))\n",
    "model.add(Activation(activation='relu', name='relu2'))\n",
    "model.add(Dense(32, name='fc3', kernel_initializer='lecun_uniform', kernel_regularizer=l1(0.0001)))\n",
    "model.add(Activation(activation='relu', name='relu3'))\n",
    "model.add(Dense(5, name='output', kernel_initializer='lecun_uniform', kernel_regularizer=l1(0.0001)))\n",
    "model.add(Activation(activation='softmax', name='softmax'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "We'll use Adam optimizer with categorical crossentropy loss.\n",
    "The callbacks will decay the learning rate and save the model into a directory 'model_1'\n",
    "The model isn't very complex, so this should just take a few minutes even on the CPU.\n",
    "If you've restarted the notebook kernel after training once, set `train = False` to load the trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = True\n",
    "if train:\n",
    "    adam = Adam(lr=0.0001)\n",
    "    model.compile(optimizer=adam, loss=['categorical_crossentropy'], metrics=['accuracy'])\n",
    "    callbacks = all_callbacks(stop_patience = 1000,\n",
    "                              lr_factor = 0.5,\n",
    "                              lr_patience = 10,\n",
    "                              lr_epsilon = 0.000001,\n",
    "                              lr_cooldown = 2,\n",
    "                              lr_minimum = 0.0000001,\n",
    "                              outputDir = 'model_1')\n",
    "    model.fit(X_train_val, y_train_val, batch_size=1024,\n",
    "              epochs=30, validation_split=0.25, shuffle=True,\n",
    "              callbacks = callbacks.callbacks)\n",
    "else:\n",
    "    from tensorflow.keras.models import load_model\n",
    "    model = load_model('model_1/KERAS_check_best_model.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check performance\n",
    "Check the accuracy and make a ROC curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotting\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import accuracy_score\n",
    "y_keras = model.predict(X_test)\n",
    "print(\"Accuracy: {}\".format(accuracy_score(np.argmax(y_test, axis=1), np.argmax(y_keras, axis=1))))\n",
    "plt.figure(figsize=(9,9))\n",
    "_ = plotting.makeRoc(y_test, y_keras, le.classes_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert the model to FPGA firmware with hls4ml\n",
    "Now we will go through the steps to convert the model we trained to a low-latency optimized FPGA firmware with hls4ml.\n",
    "First, we will evaluate its classification performance to make sure we haven't lost accuracy using the fixed-point data types. \n",
    "Then we will synthesize the model with Vivado HLS and check the metrics of latency and FPGA resource usage.\n",
    "\n",
    "## Make an hls4ml config & model\n",
    "The hls4ml Neural Network inference library is controlled through a configuration dictionary.\n",
    "In this example we'll use the most simple variation, later exercises will look at more advanced configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hls4ml\n",
    "config = hls4ml.utils.config_from_keras_model(model, granularity='model')\n",
    "print(\"-----------------------------------\")\n",
    "print(\"Configuration\")\n",
    "plotting.print_dict(config)\n",
    "print(\"-----------------------------------\")\n",
    "hls_model = hls4ml.converters.convert_from_keras_model(model,\n",
    "                                                       hls_config=config,\n",
    "                                                       output_dir='model_1/hls4ml_prj',\n",
    "                                                       fpga_part='xcu250-figd2104-2L-e')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualise what we created. The model architecture is shown, annotated with the shape and data types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls4ml.utils.plot_model(hls_model, show_shapes=True, show_precision=True, to_file=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile, predict\n",
    "Now we need to check that this model performance is still good. We compile the hls_model, and then use `hls_model.predict` to execute the FPGA firmware with bit-accurate emulation on the CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls_model.compile()\n",
    "X_test = np.ascontiguousarray(X_test)\n",
    "y_hls = hls_model.predict(X_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare\n",
    "That was easy! Now let's see how the performance compares to Keras:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Keras  Accuracy: {}\".format(accuracy_score(np.argmax(y_test, axis=1), np.argmax(y_keras, axis=1))))\n",
    "print(\"hls4ml Accuracy: {}\".format(accuracy_score(np.argmax(y_test, axis=1), np.argmax(y_hls, axis=1))))\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 9))\n",
    "_ = plotting.makeRoc(y_test, y_keras, le.classes_)\n",
    "plt.gca().set_prop_cycle(None) # reset the colors\n",
    "_ = plotting.makeRoc(y_test, y_hls, le.classes_, linestyle='--')\n",
    "\n",
    "from matplotlib.lines import Line2D\n",
    "lines = [Line2D([0], [0], ls='-'),\n",
    "         Line2D([0], [0], ls='--')]\n",
    "from matplotlib.legend import Legend\n",
    "leg = Legend(ax, lines, labels=['keras', 'hls4ml'],\n",
    "            loc='lower right', frameon=False)\n",
    "ax.add_artist(leg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthesize\n",
    "Now we'll actually use Vivado HLS to synthesize the model. We can run the build using a method of our `hls_model` object.\n",
    "After running this step, we can integrate the generated IP into a workflow to compile for a specific FPGA board.\n",
    "In this case, we'll just review the reports that Vivado HLS generates, checking the latency and resource usage.\n",
    "\n",
    "**This can take several minutes.**\n",
    "\n",
    "While the C-Synthesis is running, we can monitor the progress looking at the log file by opening a terminal from the notebook home, and executing:\n",
    "\n",
    "`tail -f model_1/hls4ml_prj/vivado_hls.log`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls_model.build(csim=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check the reports\n",
    "Print out the reports generated by Vivado HLS. Pay attention to the Latency and the 'Utilization Estimates' sections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls4ml.report.read_vivado_report('model_1/hls4ml_prj/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "Since `ReuseFactor = 1` we expect each multiplication used in the inference of our neural network to use 1 DSP. Is this what we see? (Note that the Softmax layer should use 5 DSPs, or 1 per class)\n",
    "Calculate how many multiplications are performed for the inference of this network...\n",
    "(We'll discuss the outcome)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
